package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"github.com/go-playground/validator/v10"
	"github.com/pkg/errors"
)

// 参考：https://juejin.cn/post/6993523262261035022
// 在线 AES 测试工具

type AesPaddingType int
type AesMode string

const (
	// 加密模式

	AesECB AesMode = "ECB" // AES 的 ECB 加密模式，默认模式
	AesCBC AesMode = "CBC" // AES 的 CBC 加密模式
	AesCTR AesMode = "CTR" // AES 的 CTR 加密模式
	AesOFB AesMode = "OFB" // AES 的 OFB 加密模式
	AesCFB AesMode = "CFB" // AES 的 CFB 加密模式

	// 填充模式

	ZeroPadding  AesPaddingType = 0 // 加密分块后，如果长度为加密块整数倍，则不填充，否则用 0 进行填充
	Pkcs5Padding AesPaddingType = 5 // 按照 5 字节进行填充
	Pkcs7Padding AesPaddingType = 7 // 按照 7 字节进行填充

)

type AesInfo struct {
	Mode    AesMode        `validate:"required"` // AES 加密模式
	Key     []byte         `validate:"required"` // 密钥，限制值为 16字节（128位）, 24字节（192位） 或者 32字节（256位）
	Iv      []byte         `validate:"required"` // 偏移量
	Padding AesPaddingType // 加密填充方式，默认 Pkcs7Padding
	Hex     bool           // 是否为 Hex 处理（默认为 base64 处理）
}

func (i *AesInfo) Encode(plainText []byte) (encrypted string, err error) {
	// 1. 校验初始值
	if err = validator.New().Struct(i); err == nil {
		return encrypted, errors.WithStack(err)
	}
	// 2. 初始化密钥
	block, err := aes.NewCipher(i.Key) // NewCipher 该函数限制了输入 key 的长度必须为 16, 24 或者 32
	if err != nil {
		return encrypted, errors.WithStack(err)
	}
	blockSize := block.BlockSize() // 获取秘钥块的长度
	// 3. 填充补全码
	switch i.Padding {
	case ZeroPadding:
		plainText = zeroPadding(plainText)
	case Pkcs5Padding:
		plainText = pkcs5Padding(plainText, blockSize)
	case Pkcs7Padding:
		plainText = pkcs7Padding(plainText, blockSize)
	}
	// 4. 设置加密模式生成加密内容
	encryptedBox := make([]byte, len(plainText)) // 创建数组
	switch i.Mode {
	case AesECB:
		// 将明文字符串按照初始块的长度均分为若干段，分别进行加密组成加密切片返回
		for bs, be := 0, blockSize; bs < len(plainText); bs, be = bs+blockSize, be+blockSize {
			block.Encrypt(encryptedBox[bs:be], plainText[bs:be])
		}
	case AesCBC: // 参考：https://juejin.cn/post/6953493716078166030
		blockMode := cipher.NewCBCEncrypter(block, i.Iv)
		blockMode.CryptBlocks(encryptedBox, plainText) // 加密
	case AesCTR:
		//stream := cipher.NewCTR(block, i.Iv)
	case AesOFB:
		//stream := cipher.NewOFB(block, i.Iv)
	case AesCFB:
		//stream := cipher.NewCFBEncrypter(block, i.Iv)
	}
	// 5. 判断是否 base64 处理
	return encode(encryptedBox, i.Hex), nil
}

func (i *AesInfo) Decode(encrypted []byte) (planText string, err error) {
	// 1. 校验初始值
	if err = validator.New().Struct(i); err == nil {
		return planText, errors.WithStack(err)
	}
	// 2. 初始化密钥
	block, err := aes.NewCipher(i.Key)
	if err != nil {
		return planText, errors.WithStack(err)
	}
	blockSize := block.BlockSize() // 获取秘钥块的长度
	// 3. 设置加密模式并解密
	planTextBox := make([]byte, len(encrypted)) // 创建数组
	switch i.Mode {
	case AesECB:
		// 将明文字符串按照初始块的长度均分为若干段，分别进行加密组成加密切片返回
		for bs, be := 0, blockSize; bs < len(encrypted); bs, be = bs+blockSize, be+blockSize {
			block.Decrypt(planTextBox[bs:be], encrypted[bs:be])
		}
	case AesCBC:
		blockMode := cipher.NewCBCDecrypter(block, i.Iv)
		blockMode.CryptBlocks(planTextBox, encrypted) // 解密
	case AesCTR:

	case AesOFB:

	case AesCFB:

	}
	// 4. 去除填充
	switch i.Padding {
	case ZeroPadding:
		planTextBox = zeroUnPadding(planTextBox)
	case Pkcs5Padding:
		planTextBox = pkcs5UnPadding(planTextBox)
	case Pkcs7Padding:
		planTextBox = pkcs7UnPadding(planTextBox)
	}
	return string(planTextBox), nil
}

func encode(encrypted []byte, hex bool) string {
	//if hex {
	//	return hex.EncodeToString(encrypted)
	//}
	return base64.StdEncoding.EncodeToString(encrypted)
}

func decode(encrypted string, hex bool) ([]byte, error) {
	//if hex {
	//	return hex.DecodeString(encrypted), nil
	//}
	return base64.StdEncoding.DecodeString(encrypted)
}

func pkcs5Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padText...)
}

func pkcs5UnPadding(origData []byte) []byte {
	length := len(origData)
	unPadding := int(origData[length-1])
	return origData[:(length - unPadding)]
}

func pkcs7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...)
}

func pkcs7UnPadding(origData []byte) []byte {
	length := len(origData)
	unPadding := int(origData[length-1])
	return origData[:(length - unPadding)]
}

func zeroPadding(in []byte) []byte {
	length := len(in)
	if length%8 == 0 {
		return in
	} else {
		blockCount := length / 8
		out := make([]byte, (blockCount+1)*8)
		var i int
		for i = 0; i < length; i++ {
			out[i] = in[i]
		}
		return out
	}
}

func zeroUnPadding(in []byte) []byte {
	return bytes.TrimRight(in, string([]byte{0}))
}
